package SPH;

import SPH.RadixSort.RadixSortCL;

import com.jogamp.opencl.CLBuffer;
import com.jogamp.opencl.CLCommandQueue;
import com.jogamp.opencl.CLContext;
import com.jogamp.opencl.CLDevice;
import com.jogamp.opencl.CLKernel;
import com.jogamp.opencl.CLProgram;

import java.io.File;
import java.io.IOException;
import java.nio.FloatBuffer;
import java.nio.IntBuffer;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.Dictionary;
import java.util.Hashtable;
import java.util.Iterator;
import java.util.List;
import java.util.Random;

import javax.swing.Icon;

import static java.lang.System.*;
import static com.jogamp.opencl.CLMemory.Mem.*;
import static java.lang.Math.*;

public class SPHSolver {
	
	public static Random randomGenerator = new Random();
	//TODO: I think this block we should to put to some class with constants
	public static final int PARTICLE_COUNT = ( 32 * 1024 );
	public static final int NEIGHBOR_COUNT = 32;

	static final float M_PI = 3.1415927f;
	static final int RAND_MAX = 0x7fff;

	public static final int XMIN = 0;
	public static final int XMAX = 100;
	public static final int YMIN = 0;
	public static final int YMAX = 40;
	public static final int ZMIN = 0;
	public static final int ZMAX = 40;

	public static final float xmin = XMIN;
	public static final float xmax = XMAX;
	public static final float ymin = YMIN;
	public static final float ymax = YMAX;
	public static final float zmin = ZMIN;
	public static final float zmax = ZMAX;

	public static final float rho0 = 1000.0f;
	public static final float stiffness = 0.75f;
	public static final float h = 3.34f;
	public static final float hashGridCellSize = 2.0f * h;
	public static final float hashGridCellSizeInv = 1.0f / hashGridCellSize;
	public static final float mass = 0.001f;
	public static final float simulationScale = 0.004f;
	public static final float simulationScaleInv = 1.0f / simulationScale;
	public static final float mu = 10.0f;
	public static final float timeStep = 0.0042f;
	public static final float CFLLimit = 100.0f;
	public static final int NK = NEIGHBOR_COUNT * PARTICLE_COUNT;
	public static final float damping = 0.75f;

	public static final float Wpoly6Coefficient = (float) (315.0f / ( 64.0f * M_PI * Math.pow( h * simulationScale, 9.0f ) ));
	public static final float gradWspikyCoefficient= (float) (-45.0f / ( M_PI * Math.pow( h * simulationScale, 6.0f ) ));
	public static final float del2WviscosityCoefficient = -gradWspikyCoefficient;

	public static final float gravity_x = 0.0f;
	public static final float gravity_y = -9.8f;
	public static final float gravity_z = 0.0f;
	//
	private CLContext context;
	public CLCommandQueue queue;
	private CLProgram program;
	private CLDevice device;
	public CLBuffer<FloatBuffer> acceleration;
	public CLBuffer<IntBuffer> gridCellIndex;
	public CLBuffer<IntBuffer> gridCellIndexFixedUp;
	public CLBuffer<FloatBuffer> neighborMap;
	public CLBuffer<IntBuffer> particleIndex;
	public CLBuffer<FloatBuffer> position;
	public CLBuffer<FloatBuffer> pressure;
	public CLBuffer<FloatBuffer> rho;
	public CLBuffer<FloatBuffer> rhoInv;
	public CLBuffer<FloatBuffer> sortedPosition;
	public CLBuffer<FloatBuffer> sortedVelocity;
	public CLBuffer<FloatBuffer> velocity;
	private CLKernel clearBuffers;
	private CLKernel computeAcceleration;
	private CLKernel computeDensityPressure;
	private CLKernel findNeighbors;
	private CLKernel hashParticles;
	private CLKernel indexPostPass;
	private CLKernel indexx;
	private CLKernel integrate;
	private CLKernel sortPostPass;
	
	SPHModel model;
	private RadixSortCL radixSort;
	public int gridCellsX;
	public int gridCellsY;
	public int gridCellsZ;
	public int gridCellCount;
	
	public SPHSolver()  throws IOException  {
		//SPHSolver kernelTest = new SPHSolver();
		//radixSort = new RadixSortCL();
		
		gridCellsX = (int)( ( XMAX - XMIN ) / h ) + 1;
		gridCellsY = (int)( ( YMAX - YMIN ) / h ) + 1;
		gridCellsZ = (int)( ( ZMAX - ZMIN ) / h ) + 1;
		gridCellCount = gridCellsX * gridCellsY * gridCellsZ;
		model = new SPHModel(gridCellsX, gridCellsY, gridCellsZ);
		context = CLContext.create();
		
		out.println("created "+ context);
		// an array with available devices
		CLDevice[] devices = context.getDevices();

		for(int i=0; i<devices.length; i++)
		{
			out.println("device-" + i + ": " + devices[i]);
		}	

		// have a look at the output and select a device
		device = devices[0];
		System.out.println("Version " + device.getCVersion());
		System.out.println("Version " + device.getDriverVersion());
		//System.out.println("Version" + device.get);
		// ... or use this code to select fastest device
		//CLDevice device = context.getMaxFlopsDevice();

		out.println("using "+ device);
		out.println("max workgroup size: " + device.getMaxWorkGroupSize());
		out.println("max workitems size: " + device.getMaxWorkItemSizes()[0]);
		
		// create command queue on selected device.
		queue = device.createCommandQueue();
		
		//radixSort.initializeSort(context, queue, PARTICLE_COUNT);
		
		// load sources, create and build program
		program = context.createProgram(SPHSolver.class.getResourceAsStream("/resource/sphFluidDemo.cl")).build();
		
		// NOTE: do the same for all the neurons, C. elegans has 302 neurons, even if they don't fire and this is HH (squid) it will hopefully give a first indication
		int elementCount = PARTICLE_COUNT;                                  // Length of arrays to process
		//int localWorkSize = min((int)kernel.getWorkGroupSize(device), 256);  // Local work size dimensions
		int globalWorkSize = elementCount * 4;//roundUp(localWorkSize, elementCount);   // rounded up to the nearest multiple of the localWorkSize
		//int globalWorkSize_Plotting = roundUp(localWorkSize, elementCount*steps);
		/* input buffers declarations */
		acceleration = context.createFloatBuffer(PARTICLE_COUNT * 4, READ_WRITE);
		gridCellIndex = context.createIntBuffer(gridCellCount + 1, READ_WRITE);
		gridCellIndexFixedUp = context.createIntBuffer(gridCellCount + 1, READ_WRITE);
		neighborMap = context.createFloatBuffer(NK * 2, READ_WRITE);
		particleIndex = context.createIntBuffer(PARTICLE_COUNT * 2, READ_WRITE);
		position = context.createFloatBuffer(PARTICLE_COUNT * 4, READ_WRITE);
		pressure = context.createFloatBuffer(PARTICLE_COUNT * 1, READ_WRITE);
		rho = context.createFloatBuffer(PARTICLE_COUNT * 1, READ_WRITE);
		rhoInv = context.createFloatBuffer(PARTICLE_COUNT * 1, READ_WRITE);
		sortedPosition = context.createFloatBuffer(PARTICLE_COUNT * 4, READ_WRITE);
		sortedVelocity = context.createFloatBuffer(PARTICLE_COUNT * 4, READ_WRITE);
		velocity = context.createFloatBuffer(PARTICLE_COUNT * 4, READ_WRITE);
		
		/*kernels*/
		clearBuffers = program.createCLKernel("clearBuffers");
		computeAcceleration = program.createCLKernel("computeAcceleration");
		computeDensityPressure = program.createCLKernel("computeDensityPressure");
		findNeighbors = program.createCLKernel("findNeighbors");
		hashParticles = program.createCLKernel("hashParticles");
		indexPostPass = program.createCLKernel("indexPostPass");
		indexx = program.createCLKernel("indexx");
		integrate = program.createCLKernel("integrate");
		sortPostPass = program.createCLKernel("sortPostPass");
		SecureRandom random = new SecureRandom();
		for(int i = 0;i<PARTICLE_COUNT;i++){
			float r1 = ((float)randomGenerator.nextInt(RAND_MAX) / (float)RAND_MAX );//random.nextFloat();
			float r = ((float)randomGenerator.nextInt(RAND_MAX) / (float)RAND_MAX );//tr.rand.next();
			Vector3D positionVector = new Vector3D();
			Vector3D velocityVector = new Vector3D();
			positionVector.x = scale(XMIN, XMAX/10 , r); 
			r = ((float)randomGenerator.nextInt(RAND_MAX) / (float)RAND_MAX );//tr.rand.next();
			positionVector.y = scale(YMIN, YMAX , r); 
			r = ((float)randomGenerator.nextInt(RAND_MAX) / (float)RAND_MAX );//tr.rand.next();
			positionVector.z = scale(ZMIN, ZMAX , r);
			positionVector.p = 0;
			if(position.getBuffer().remaining() != 0){
				position.getBuffer().put(positionVector.x);
				position.getBuffer().put(positionVector.y);
				position.getBuffer().put(positionVector.z);
				position.getBuffer().put(positionVector.p);
			}	
			r = ((float)randomGenerator.nextInt(RAND_MAX) / (float)RAND_MAX );//tr.rand.next();
			velocityVector.x = scale(-1.0f, 1.0f, r);
			r = ((float)randomGenerator.nextInt(RAND_MAX) / (float)RAND_MAX );//tr.rand.next();
			velocityVector.y = scale(-1.0f, 1.0f, r);
			r = ((float)randomGenerator.nextInt(RAND_MAX) / (float)RAND_MAX );//tr.rand.next();
			velocityVector.z = scale(-1.0f, 1.0f, r);
			velocityVector.p = 0;
			if(velocity.getBuffer().remaining() != 0){
				velocity.getBuffer().put(velocityVector.x);
				velocity.getBuffer().put(velocityVector.y);
				velocity.getBuffer().put(velocityVector.z);
				velocity.getBuffer().put(velocityVector.p);
			}
			SPHParticle particle = new SPHParticle(positionVector, velocityVector, 1);
			model.particles.add(particle);
			
		}
		position.getBuffer().rewind();
		velocity.getBuffer().rewind();
		queue.putWriteBuffer(position, false)
		.putWriteBuffer(velocity, false);
		queue.finish();
	}
	
	public void cleanContext(){
		this.context.release();
	}
	
	public int _runClearBuffers(){
		int elementCount = PARTICLE_COUNT;                                  // Length of arrays to process
		int localWorkSize = min((int)clearBuffers.getWorkGroupSize(device), 256);  // Local work size dimensions
		int globalWorkSize = roundUp(localWorkSize, elementCount);   // rounded up to the nearest multiple of the localWorkSize
		clearBuffers.setArg(0, neighborMap);
		queue.put1DRangeKernel(clearBuffers, 0, globalWorkSize, localWorkSize);
		return 0;
	}
	
	public int _runComputeAcceleration(){
		computeAcceleration.setArg( 0, neighborMap );
		computeAcceleration.setArg( 1, pressure );
		computeAcceleration.setArg( 2, rho );
		computeAcceleration.setArg( 3, rhoInv );
		computeAcceleration.setArg( 4, sortedPosition );
		computeAcceleration.setArg( 5, sortedVelocity );
		computeAcceleration.setArg( 6, CFLLimit );
		computeAcceleration.setArg( 7, del2WviscosityCoefficient );
		computeAcceleration.setArg( 8, gradWspikyCoefficient );
		computeAcceleration.setArg( 9, h );
		computeAcceleration.setArg( 10, mass );
		computeAcceleration.setArg( 11, mu );
		computeAcceleration.setArg( 12, simulationScale );
		computeAcceleration.setArg( 13, acceleration );
		//int elementCount = PARTICLE_COUNT;                                  // Length of arrays to process
		//int localWorkSize = min((int)computeAcceleration.getWorkGroupSize(device), 256);  // Local work size dimensions
		//int globalWorkSize = roundUp(localWorkSize, elementCount);   // rounded up to the nearest multiple of the localWorkSize
		queue.put1DRangeKernel(computeAcceleration, 0, PARTICLE_COUNT, 256);
		return 0;
	}
	
	public int _runComputeDensityPressure(){
		computeDensityPressure.setArg( 0, neighborMap );
		computeDensityPressure.setArg( 1, Wpoly6Coefficient );
		computeDensityPressure.setArg( 2, h );
		computeDensityPressure.setArg( 3, mass );
		computeDensityPressure.setArg( 4, rho0 );
		computeDensityPressure.setArg( 5, simulationScale );
		computeDensityPressure.setArg( 6, stiffness );
		computeDensityPressure.setArg( 7, pressure );
		computeDensityPressure.setArg( 8, rho );
		computeDensityPressure.setArg( 9, rhoInv );
		int elementCount = PARTICLE_COUNT;                                  // Length of arrays to process
		int localWorkSize = min((int)computeDensityPressure.getWorkGroupSize(device), 256);  // Local work size dimensions
		int globalWorkSize = roundUp(localWorkSize, elementCount);   // rounded up to the nearest multiple of the localWorkSize
		queue.put1DRangeKernel(computeDensityPressure, 0, globalWorkSize, localWorkSize);
		return 0;
	}
	
	public int _runFindNeighbors(){
		findNeighbors.setArg( 0, gridCellIndexFixedUp );
		findNeighbors.setArg( 1, sortedPosition );
		gridCellCount = ((gridCellsX) * (gridCellsY)) * (gridCellsZ);
		findNeighbors.setArg( 2, gridCellCount );
		findNeighbors.setArg( 3, gridCellsX );
		findNeighbors.setArg( 4, gridCellsY );
		findNeighbors.setArg( 5, gridCellsZ );
		findNeighbors.setArg( 6, h );
		findNeighbors.setArg( 7, hashGridCellSize );
		findNeighbors.setArg( 8, hashGridCellSizeInv );
		findNeighbors.setArg( 9, simulationScale );
		findNeighbors.setArg( 10, xmin );
		findNeighbors.setArg( 11, ymin );
		findNeighbors.setArg( 12, zmin );
		findNeighbors.setArg( 13, neighborMap );
		int elementCount = PARTICLE_COUNT;                                  // Length of arrays to process
		int localWorkSize = min((int)hashParticles.getWorkGroupSize(device), 256);  // Local work size dimensions
		int globalWorkSize = roundUp(localWorkSize, elementCount);   // rounded up to the nearest multiple of the localWorkSize
		queue.put1DRangeKernel(findNeighbors, 0, globalWorkSize, localWorkSize);
		return 0;
	}
	
	public int _runHashParticles(){
		// Stage HashParticles
		hashParticles.setArg( 0, position );
		hashParticles.setArg( 1, gridCellsX );
		hashParticles.setArg( 2, gridCellsY );
		hashParticles.setArg( 3, gridCellsZ );
		hashParticles.setArg( 4, hashGridCellSizeInv );
		hashParticles.setArg( 5, xmin );
		hashParticles.setArg( 6, ymin );
		hashParticles.setArg( 7, zmin );
		hashParticles.setArg( 8, particleIndex );
		int elementCount = PARTICLE_COUNT;                                  // Length of arrays to process
		int localWorkSize = 256;//min((int)hashParticles.getWorkGroupSize(device), 256);  // Local work size dimensions
		int globalWorkSize = PARTICLE_COUNT;//roundUp(localWorkSize, elementCount);   // rounded up to the nearest multiple of the localWorkSize
		queue.put1DRangeKernel(hashParticles, 0, globalWorkSize, localWorkSize);
		return 0;
	}
	
	public int _runIndexPostPass(){
		indexPostPass.setArg( 0, gridCellIndex );
		gridCellCount = ((gridCellsX) * (gridCellsY)) * (gridCellsZ);
		indexPostPass.setArg( 1, gridCellCount );
		indexPostPass.setArg( 2, gridCellIndexFixedUp );
		int gridCellCountRoundedUp = ((( gridCellCount - 1 ) / 256 ) + 1 ) * 256;
		int localWorkSize = 256;//min((int)indexPostPass.getWorkGroupSize(device), 256);  // Local work size dimensions
		//int globalWorkSize = roundUp(localWorkSize, gridCellCountRoundedUp);   // rounded up to the nearest multiple of the localWorkSize
		queue.put1DRangeKernel(indexPostPass, 0, gridCellCountRoundedUp, localWorkSize);
		return 0;
	}
	
	public int _runIndexx(){
		// Stage Indexx
		indexx.setArg( 0, particleIndex );
		gridCellCount = ((gridCellsX) * (gridCellsY)) * (gridCellsZ);
		indexx.setArg( 1, gridCellCount );
		indexx.setArg( 2, gridCellIndex );
		int gridCellCountRoundedUp = ((( gridCellCount - 1 ) / 256 ) + 1 ) * 256;
		int localWorkSize = min((int)indexx.getWorkGroupSize(device), 256);  // Local work size dimensions
		queue.put1DRangeKernel(indexx, 0, gridCellCountRoundedUp, localWorkSize);
		return 0;
	}
	
	public int _runIntegrate(){
		// Stage Integrate
		integrate.setArg( 0, acceleration );
		integrate.setArg( 1, sortedPosition );
		integrate.setArg( 2, sortedVelocity );
		integrate.setArg( 3, gravity_x );
		integrate.setArg( 4, gravity_y );
		integrate.setArg( 5, gravity_z );
		integrate.setArg( 6, simulationScaleInv );
		integrate.setArg( 7, timeStep );
		integrate.setArg( 8, xmin );
		integrate.setArg( 9, xmax );
		integrate.setArg( 10, ymin );
		integrate.setArg( 11, ymax );
		integrate.setArg( 12, zmin );
		integrate.setArg( 13, zmax );
		integrate.setArg( 14, damping );
		integrate.setArg( 15, position );
		integrate.setArg( 16, velocity );
		int elementCount = PARTICLE_COUNT;                                  // Length of arrays to process
		int localWorkSize = min((int)integrate.getWorkGroupSize(device), 256);  // Local work size dimensions
		int globalWorkSize = roundUp(localWorkSize, elementCount);   // rounded up to the nearest multiple of the localWorkSize
		queue.put1DRangeKernel(integrate, 0, globalWorkSize, localWorkSize);
		return 0;
	}
	
	public int _runSortPostPass(){
		// Stage SortPostPass
		sortPostPass.setArg( 0, particleIndex );
		sortPostPass.setArg( 1, position );
		sortPostPass.setArg( 2, velocity );
		sortPostPass.setArg( 3, sortedPosition );
		sortPostPass.setArg( 4, sortedVelocity );
		int elementCount = PARTICLE_COUNT;                                  // Length of arrays to process
		int localWorkSize = 256;//min((int)sortPostPass.getWorkGroupSize(device), 256);  // Local work size dimensions
		int globalWorkSize = PARTICLE_COUNT;//roundUp(localWorkSize, elementCount);   // rounded up to the nearest multiple of the localWorkSize
		queue.put1DRangeKernel(sortPostPass, 0, globalWorkSize, localWorkSize);
		return 0;
	}
	
	
	public int _runSort(){
		//radixSort.sort( particleIndex );
		//radixSort.wait();
		
		//this version for debug mode
		List<int[]> _particleIndex = new ArrayList<int[]>();//new int[PARTICLE_COUNT*2];
		queue.putReadBuffer(particleIndex, true);
		queue.finish();
		for(int i = 0; i < PARTICLE_COUNT * 2;i+=2){
			int[] element = {particleIndex.getBuffer().get(i), particleIndex.getBuffer().get(i + 1)};
			_particleIndex.add(element);
		}
		Collections.sort(_particleIndex, new MyCompare());
		
		//Arrays.sort(_particleIndex,0,PARTICLE_COUNT, new MyCompare());//   (_particleIndex);
		for(int i = 0; i< _particleIndex.size();i++){
			if(particleIndex.getBuffer().remaining() != 0){
				for(int j=0;j<2;j++)
					particleIndex.getBuffer().put(_particleIndex.get(i)[j]);
			}
		}
		particleIndex.getBuffer().rewind();
		queue.putWriteBuffer(particleIndex, true);
		queue.finish();	
		return 0;
	}
	
	class MyCompare implements Comparator<int[]>{
		@Override
		public int compare(int[] o1, int[] o2) {
			// TODO Auto-generated method stub
			if( o1[0] < o2[0] ) return -1;
			if( o1[0] > o2[0]) return +1;
			return 0;
		}
	}
	void advanceInTime(){
		step();
		for (int i = 0;i<model.particles.size();i++) {
			model.particles.get(i).positionVector.x = position.getBuffer().get(4*i);
			model.particles.get(i).positionVector.y = position.getBuffer().get(4*i+1);
			model.particles.get(i).positionVector.z = position.getBuffer().get(4*i+2);
			model.particles.get(i).positionVector.p = position.getBuffer().get(4*i+3);
		}
	}
	
	private void step(){
		_runClearBuffers();
		_runHashParticles();
		_runSort();
		_runSortPostPass();
		_runIndexx();
		_runIndexPostPass();
		_runFindNeighbors();
		_runComputeDensityPressure();
		_runComputeAcceleration();
		_runIntegrate();
		queue.putReadBuffer(position, true);
		queue.finish();
	}
	
	private static float scale(float min, float max, float r){
		return (float)(min + r * (max - min));
	}
	
	private static int roundUp(int groupSize, int globalSize) {
        int r = globalSize % groupSize;
        if (r == 0) {
            return globalSize;
        } else {
            return globalSize + groupSize - r;
        }
    }
	
	private void initializeOpenCl(){
		
	}

	public void finishQueue() {
		// TODO Auto-generated method stub
		queue.finish();
	}
}
